# -*- coding: utf-8 -*-
"""
优化检索结果方案：
(1) 多路召回后，使用rrf算法，手动选出topk
(2) 使用reranker模型，对多路召回的chunks,进行排序选出topk
"""
import re
from typing import List

from langchain_core.output_parsers import StrOutputParser
from langchain.prompts import PromptTemplate
from langchain.schema.runnable import RunnablePassthrough
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import xml.etree.ElementTree as et
from lxml.html import XHTMLParser
from ycr.Ai_Chat_Code_main.xunfei.spark_model import WrappedSparkModel
from loguru import logger as log


def rewrite_question(question, history=None):
    '''
    问题的改写
    :param question:
    :param history: chat history
    :return:
    '''
    template = '''
    如果你是一位问题改写专家，请对以下问题进行问题的拆解改写
    改写要求：
        1.先分析问题，对原始问题进行独立语义拆解
        2.拆解后，对每个拆解语义进行问题改写，每一个改写后的问题只包含一个语义独立问题
        3.确保改写后的问题依然清晰易懂
    注意：具体改写的流程严格按照以下例子,直接输出改写后的问题，按照xml格式输出
    以下是一些例子:
    examples:
        example1(最常见，最简单的例子):
            question: 中国邮政储蓄银行合肥基地（四期）建设项目的招标截止时间？
            history: xxxx
            workflow:
                让我们一步一步来分析
                (1)首先对问题进行拆解，此问题本身只有一个独立的语义，
                    定语是“中国邮政储蓄银行合肥基地（四期）建设项目”，问题是“招标截止时间”
                (2)对拆解结果进行问题重写：“中国邮政储蓄银行合肥基地（四期）建设项目 招标何时截止？”
                (3)对(2)中的重写问题进行结构化输出，例如使用xml进行封装，最终答案如下：
                    <rewrite_questions>
                        <question_1>中国邮政储蓄银行合肥基地（四期）建设项目 招标何时截止？</question_1>
                    </rewrite_questions>
        example2(缺少限定语的问题):
            question:这个项目的招标人是谁？
            history: ....xxx项目....
            workflow:
                让我们一步一步来分析
                (1)首先对问题进行拆解，此问题本身只有一个独立语义问题，但是找不到这个项目具体是哪个？
                    所以尝试从history中找项目名称，如果找到了具体的项目名称，定语是“xxx项目”，问题是“招标人是谁？”
                (2)对拆解结果进行问题重写：“xxx项目招标人是谁？”
                (3)对(2)中的重写问题进行结构化输出，例如使用xml进行封装，最终答案如下：
                    <rewrite_questions>
                        <question_1>xxx项目招标人是谁？</question_1>
                    </rewrite_questions>
         example3(如果存在多个独立语义的情况):
            question:合肥市第七中学宿舍综合楼新建及校园维修项目的资金来源和项目出资比例是多少？
            workflow:
                让我们一步一步来分析
                (1)首先对问题进行拆解，原始问题包含两个独立语义问题:
                    定语是“合肥市第七中学宿舍综合楼新建及校园维修项目”，独立语义问题包含两个
                        "问题1":"资金来源",
                        "问题2":"项目出资比例"
                (2)对拆解结果进行问题重写：
                    问题1重写：合肥市第七中学宿舍综合楼新建及校园维修项目资金来源？
                    问题2重写：合肥市第七中学宿舍综合楼新建及校园维修项目项目出资比例？
                (3)对(2)中的重写问题进行结构化输出，例如使用xml进行封装，最终答案如下：
                    <rewrite_questions>
                        <question_1>合肥市第七中学宿舍综合楼新建及校园维修项目资金来源？</question_1>
                        <question_2>合肥市第七中学宿舍综合楼新建及校园维修项目项目出资比例？</question_2>
                    </rewrite_questions>
            
        以下是实际要拆解改写的问题:
        {question}
        '''
    rag_prompt = PromptTemplate.from_template(template=template)  # 创建提示词模版
    rag_chain = (
            rag_prompt
            | WrappedSparkModel.get_spark_chat_max_model()
            | StrOutputParser()
    )
    rewrite_questions = rag_chain.invoke({"question": question, "history":history})
    log.info('rewrite_questions={}', rewrite_questions)
    # 将问题转化为列表返回
    rewrite_questions_list = None
    if rewrite_questions:
        try:
            root = et.fromstring(rewrite_questions.strip(), parser=XHTMLParser())
            rewrite_questions_list = root.xpath('//rewrite_questions//text()[normalize-space()]')
        except Exception as e:
            log.error("rewrite_question has an error={}", e)
            # 使用正则简单处理
            patten_obj = re.compile(r'<question_\d+>(.*?)</question_\d+>', re.S)
            rewrite_questions_list = patten_obj.findall(rewrite_questions)
    log.info("rewrite_question rewrite_questions_list={}", rewrite_questions_list)
    return rewrite_questions_list


def get_rewrite_questions(origin_question) -> List[str]:
    """
     获取重写的问题
     param:origin_question:原始问题
     return:改写后的问题list
    """
    # 1.通过llm重写问题
    rewrite_questions = rewrite_question(origin_question)
    # 2.架构问题



def reranker(qa_pairs: List[List]):
    '''
    重排
    :param qa_pairs: question answer 问答对,注意保证传入的问答对是去重之后的
    :return:
    '''
    tokenizer = AutoTokenizer.from_pretrained(r'model/bge-reranker-v2-m3')
    model = AutoModelForSequenceClassification.from_pretrained(r'model/bge-reranker-v2-m3')
    model.eval()
    # pairs = [["临湖社区玉屏南路商业街区综合治理工程项目设计的招标条件是什么？", """
    # "1. 招标条件
    #
    #  1.1 项目名称：临湖社区玉屏南路商业街区综合治理工程项目设计
    #
    #  1.2 项目审批、核准或备案机关名称：合肥经济技术开发区经济发展局
    #
    #  1.3 批文名称及编号：《关于临湖社区玉屏南路商业街区综合治理工程项目立项的批复》，合经投〔2024〕41号
    #
    #  1.4 招标人：合肥经济技术开发区重点工程建设管理中心
    #
    #  1.5 项目业主：合肥经济技术开发区重点工程建设管理中心
    #
    #  1.6 资金来源：政府投资
    #
    #  1.7 项目出资比例：100%
    #
    #  1.8 资金落实情况: 已落实
    #
    # """], ['临湖社区玉屏南路商业街区综合治理工程项目设计的招标条件是什么？',
    #        "1.1 项目名称：临湖社区玉屏南路商业街区综合治理工程项目设计"]]
    pairs = qa_pairs
    with torch.no_grad():
        inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
        scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
        log.info("scores={}", scores)
        list_version = scores.tolist()
        # sorted_data = sorted(list_version)
        sorted_data = sorted(enumerate(list_version), key=lambda x: x[1], reverse=True)
        sorted_indices = [x[0] for x in sorted_data]
        return sorted_indices


def reciprocal_rank_fusion(document_ranks, k=60):
    """
      多路召回后，使用rrf算法，手动选出topk
      返回chunk是去重后的
    """
    fused_scores = {}
    for docs in document_ranks:
        for rank, doc in enumerate(docs):
            # 暂时简单处理
            doc_str = doc
            if doc_str not in fused_scores:
                fused_scores[doc_str] = 0
            fused_scores[doc_str] += 1 / (rank + k)
    reranked_results = [
        # 加载具体文档操作
        # (loads(doc), score)
        # 暂时简单处理
        (doc, score)
        for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    ]
    return reranked_results


if __name__ == '__main__':
    # reranker("", "")
    rewrite_question("临湖社区玉屏南路商业街区综合治理工程项目设计的投标人条件以及资金来源？")
    # rewrite_questions="""
    #  <rewrite_questions>
    #                     <question_1>合肥市第七中学宿舍综合楼新建及校园维修项目资金来源？</question_1>
    #                     <question_2>合肥市第七中学宿舍综合楼新建及校园维修项目项目出资比例？</question_2>
    #                 </rewrite_questions>
    # """
    # patten_obj = re.compile(r'<question_\d+>(.*?)</question_\d+>', re.S)
    # rewrite_questions_list = patten_obj.findall(rewrite_questions)
    # print(rewrite_questions_list)